//
// Created by apple on 2022/10/4.
//

#ifndef COROUTINE_CTX_COROUTINE_H
#define COROUTINE_CTX_COROUTINE_H

#include <ctime>
#include <cstdint>

/*
 * Abstract Class for Coroutine, its status, and utils
 * Using by the two implementations
 */

enum co_status {
    NEW = 0,
    RUNNING,
    WAITING,
    DEAD
};

template<typename T>
class AbsRoutine {
public:
    virtual T resume() = 0;

    co_status status() { return _status; }

protected:
    int resume_num = 0;
    co_status _status = NEW;
};

uint64_t timeStampMs() {
    timespec ts{};
    clock_gettime(CLOCK_MONOTONIC, &ts);
    return ts.tv_nsec / 1000000 + ts.tv_sec * 1000ull;
}

/***************************************************************
 * @brief The macro for switch and struct based coroutine implementation
 *
 * @Advantages：
 * 1.   Simple to implement, no external dependencies
 * 2.   No assembly code involved, easy to cross-platform
 * 3.   Support arbitrary parameters and return values
 * 4.   No need to manually manage the stack of functions
 * 5.   Support Nest Coroutine
 *
 * @Disadvantages:
 *  1.  This kind of implementation can only support Non-Preemptive Coroutine,
 *      because switch can only jump to the pre-set location not to an arbitrary
 *  2.  All variables must be declared as members of the class
 *      in order to achieve context preservation
 *
 * @Sample You can use these macro like here:

    class range: public AbsRoutine<int>{
    public:
        explicit range(int max): max(max) {}
        int resume() override{
            CO_BEGIN();
                    for(i = 0; i < max - 1; ++i){
                        YIELD(i);
                    }
            CO_END(i);
        }
    private:
        int max{}, i=0;
    };

    int main(){
      for(auto r = range(10); r.status() != DEAD;)
          printf("%d ", r.resume());
      return 0;
    }
 **************************************************************/

#define CO_BEGIN()  _status=RUNNING; switch(resume_num) { case 0:
#define YIELD(x)    resume_num = __LINE__; _status=WAITING; return x; case __LINE__:
#define CO_END(x)   default:{_status=DEAD; return x;}}



/***************************************************************
 * @brief The Asm based coroutine implementation
 *
 * @note Place the code under the namespace co_asm
 *
 * @Advantages
 * 1.   It can support Preemptive Coroutine
 *      (But only the non-preemptive version is implemented here)
 * 2.   You can declare any variable in coroutine functions,
 *      just like other normal ones.
 *
 * @Disadvantages:
 * 1.   Only Support x86-64 platform with Linux or Mac System.
 *      And must open O2 or stronger optimization on Linux(g++)
 * 2.   Only functions with no parameters and no return value are supported
 * 3.   Need to manually manage the stack of functions
 * 4.   Not support Nest Coroutines
 *      (You can not get back to the original entry if using nest coroutine)
 **************************************************************/


#include <unistd.h>
#include <iostream>
#include <queue>
#include <vector>
#include <cassert>

namespace co_asm {

    extern "C" void swap_context(void *, void *) asm("swap_context");
    /***************************************************************
     * @brief     save current context into the first struct, and load the second one
     **************************************************************/

    asm(R"(
swap_context:
  mov 0x00(%rsp), %rdx
  lea 0x08(%rsp), %rcx
  mov %r12, 0x00(%rdi)
  mov %r13, 0x08(%rdi)
  mov %r14, 0x10(%rdi)
  mov %r15, 0x18(%rdi)
  mov %rdx, 0x20(%rdi)
  mov %rcx, 0x28(%rdi)
  mov %rbx, 0x30(%rdi)
  mov %rbp, 0x38(%rdi)
  mov 0x00(%rsi), %r12
  mov 0x08(%rsi), %r13
  mov 0x10(%rsi), %r14
  mov 0x18(%rsi), %r15
  mov 0x20(%rsi), %rax
  mov 0x28(%rsi), %rcx
  mov 0x30(%rsi), %rbx
  mov 0x38(%rsi), %rbp
  mov %rcx, %rsp
  jmpq *%rax
)");

    struct Context {
        /***************************************************************
         * @brief   The Context of a Coroutine.
         *          Store the register value and the stack
         **************************************************************/
        void *reg[8]{};
        std::vector<char> mem;

        explicit Context(void (*func)() = nullptr) : mem(4096) {
            reg[4] = (void *) func;
            reg[5] = (char *) ((uintptr_t) (&mem.back()) & ~15ull) - sizeof(void *);
        }
    } ma;

    co_status cur_status = NEW;

    class AsmRoutine : public AbsRoutine<co_status> {
    public:
        explicit AsmRoutine(void (*func)()) : ctx(Context{func}) {}

        co_status resume() override;

        inline Context *context() { return &ctx; }

    private:
        Context ctx;

    };

    AsmRoutine *current = nullptr;

    co_status AsmRoutine::resume() {
        if (_status == DEAD)
            return _status;

        _status = RUNNING;
        current = this;
        swap_context(&ma, current->context());

        assert(cur_status == WAITING || cur_status == DEAD);
        _status = cur_status;
        return _status;
    }


    struct Task {
        /***************************************************************
         * @brief   To store expired sleeping time (ms) and the corresponding coroutine
         **************************************************************/
        uint64_t expire;
        AsmRoutine *coroutine;

        // Define the cmp. To make the priority queue a min heap.
        bool operator<(const Task &other) const { return expire > other.expire; }
    };

    std::priority_queue<Task> tasks; // using the expired sleep time as priority

    /*
     * Coroutine Event
     * Switching Context using global variables
     */

    void sleep(int ms) {
        /***************************************************************
         * @brief   Sleep and Yield. Turn the status to WAITING
         * @param   ms  Time to sleep in milliseconds
         **************************************************************/
        tasks.push(Task{.expire=timeStampMs() + ms, .coroutine=current});
        cur_status = WAITING;
        swap_context(current->context(), &ma);
    }

    inline void yield() { sleep(0); }

    void end() {
        /***************************************************************
         * @brief   The end of a coroutine. Turn the status to DEAD
         **************************************************************/
        cur_status = DEAD;
        tasks.push(Task{.expire = timeStampMs(), .coroutine = current});
        swap_context(current->context(), &ma);
    }

    void event_loop(int timeout_in_seconds) {
        /***************************************************************
         * @brief   Schedule all coroutines in tasks using priority queue
         * @param   timeout_in_seconds  Timeout limit in seconds.
         *                              Return when timeout, no matter the queue is empty or not.
         **************************************************************/
        uint64_t start = timeStampMs();
        while (true) {
            while (!tasks.empty()) {
                if (timeStampMs() > tasks.top().expire) {
                    auto task = tasks.top();
                    tasks.pop();
                    task.coroutine->resume();
                } else
                    break;
            }
            if (tasks.empty() || (timeStampMs() - start) > timeout_in_seconds * 1000)
                break;
            usleep(100);
        }
    }

    void event_loop(std::vector<AsmRoutine> &task, int timeout_in_seconds) {
        /***************************************************************
         * @brief   Resume all coroutines in task once first, and schedule them using priority queue
         * @param   timeout_in_seconds  Timeout limit in seconds.
         *                              Return when timeout, no matter the queue is empty or not.
         **************************************************************/
        for (auto &t: task)
            t.resume();
        event_loop(timeout_in_seconds);
    }

}
#endif //COROUTINE_CTX_COROUTINE_H
